# -*- coding: utf-8 -*-

import numpy as np
import tensorflow as tf
from read_data import get_images_path, process_image


def convolution(in_put, kernel_height, kernel_width,
                in_channel, out_channel, layer_name):
    with tf.variable_scope(layer_name):
        weights = tf.get_variable(name="weights",
                                  shape=[kernel_height, kernel_width, in_channel, out_channel],
                                  dtype=tf.float32,
                                  initializer=tf.contrib.layers.xavier_initializer_conv2d(),
                                  regularizer=tf.contrib.layers.l2_regularizer(0.00005))
        biases = tf.get_variable(name="biases",
                                 shape=[out_channel],
                                 dtype=tf.float32,
                                 initializer=tf.constant_initializer())
        convolution_output = tf.nn.conv2d(input=in_put,
                                          filter=weights,
                                          strides=[1, 1, 1, 1],
                                          padding="SAME",
                                          name="convolution")
        output = tf.nn.relu(tf.nn.bias_add(convolution_output, biases), name="relu")
    return output


def cnn_column_1(in_put):
    with tf.variable_scope("cnn_column_1"):
        layer_1 = convolution(in_put=in_put, kernel_height=5, kernel_width=5,
                              in_channel=3, out_channel=24, layer_name="layer_1")
        pool_1 = tf.nn.max_pool(value=layer_1, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1],
                                padding="SAME")
        layer_2 = convolution(in_put=pool_1, kernel_height=3, kernel_width=3,
                              in_channel=24, out_channel=48, layer_name="layer_2")
        pool_2 = tf.nn.max_pool(value=layer_2, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1],
                                padding="SAME")
        layer_3 = convolution(in_put=pool_2, kernel_height=3, kernel_width=3,
                              in_channel=48, out_channel=24, layer_name="layer_3")
        layer_4 = convolution(in_put=layer_3, kernel_height=3, kernel_width=3,
                              in_channel=24, out_channel=12, layer_name="layer_4")
        with tf.variable_scope("layer_5"):
            weights = tf.get_variable(name="weights",
                                      shape=[1, 1, 12, 1],
                                      dtype=tf.float32,
                                      initializer=tf.contrib.layers.xavier_initializer_conv2d(),
                                      regularizer=tf.contrib.layers.l2_regularizer(0.00005))
            layer_5 = tf.nn.conv2d(input=layer_4,
                                   filter=weights,
                                   strides=[1, 1, 1, 1],
                                   padding="SAME",
                                   name="convolution")
        # layer_5 = convolution(in_put=layer_4, kernel_height=1, kernel_width=1,
        #                       in_channel=12, out_channel=1, layer_name="layer_5")
        return layer_5


def cnn_column_2(in_put):
    with tf.variable_scope("cnn_column_1"):
        layer_1 = convolution(in_put=in_put, kernel_height=7, kernel_width=7,
                              in_channel=3, out_channel=20, layer_name="layer_1")
        pool_1 = tf.nn.max_pool(value=layer_1, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1],
                                padding="SAME")
        layer_2 = convolution(in_put=pool_1, kernel_height=5, kernel_width=5,
                              in_channel=20, out_channel=40, layer_name="layer_2")
        pool_2 = tf.nn.max_pool(value=layer_2, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1],
                                padding="SAME")
        layer_3 = convolution(in_put=pool_2, kernel_height=5, kernel_width=5,
                              in_channel=40, out_channel=20, layer_name="layer_3")
        layer_4 = convolution(in_put=layer_3, kernel_height=5, kernel_width=5,
                              in_channel=20, out_channel=10, layer_name="layer_4")
        with tf.variable_scope("layer_5"):
            weights = tf.get_variable(name="weights",
                                      shape=[1, 1, 10, 1],
                                      dtype=tf.float32,
                                      initializer=tf.contrib.layers.xavier_initializer_conv2d(),
                                      regularizer=tf.contrib.layers.l2_regularizer(0.00001))
            layer_5 = tf.nn.conv2d(input=layer_4,
                                   filter=weights,
                                   strides=[1, 1, 1, 1],
                                   padding="SAME",
                                   name="convolution")
        return layer_5


def cnn_column_3(in_put):
    with tf.variable_scope("cnn_column_1"):
        layer_1 = convolution(in_put=in_put, kernel_height=9, kernel_width=9,
                              in_channel=3, out_channel=16, layer_name="layer_1")
        pool_1 = tf.nn.max_pool(value=layer_1, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1],
                                padding="SAME")
        layer_2 = convolution(in_put=pool_1, kernel_height=7, kernel_width=7,
                              in_channel=16, out_channel=32, layer_name="layer_2")
        pool_2 = tf.nn.max_pool(value=layer_2, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1],
                                padding="SAME")
        layer_3 = convolution(in_put=pool_2, kernel_height=7, kernel_width=7,
                              in_channel=32, out_channel=16, layer_name="layer_3")
        layer_4 = convolution(in_put=layer_3, kernel_height=7, kernel_width=7,
                              in_channel=16, out_channel=8, layer_name="layer_4")
        with tf.variable_scope("layer_5"):
            weights = tf.get_variable(name="weights",
                                      shape=[1, 1, 8, 1],
                                      dtype=tf.float32,
                                      initializer=tf.contrib.layers.xavier_initializer_conv2d(),
                                      regularizer=tf.contrib.layers.l2_regularizer(0.00005))
            layer_5 = tf.nn.conv2d(input=layer_4,
                                   filter=weights,
                                   strides=[1, 1, 1, 1],
                                   padding="SAME",
                                   name="convolution")
        return layer_5


def main():
    folder = "/home/lijun/Dataset/crowd_counting/part_a/train"
    images_path, ground_truths_path = get_images_path(folder)
    with tf.Graph().as_default():
        in_put_x = tf.placeholder(dtype=tf.float32, shape=[1, None, None, 3])
        in_put_y = tf.placeholder(dtype=tf.float32, shape=[1, None, None, 1])
        predict_op = cnn_column_1(in_put=in_put_x)
        inference_loss_op = tf.nn.l2_loss(tf.sub(in_put_y, predict_op))
        with tf.Session() as session:
            session.run(tf.initialize_all_variables())
            for image_path, ground_truth_path in zip(images_path, ground_truths_path):
                image, ground_truth = process_image(image_path, ground_truth_path)
                image = np.expand_dims(image, axis=0)
                ground_truth = np.expand_dims(np.expand_dims(ground_truth, axis=0), axis=-1)
                print(image.shape, ground_truth.shape)
                inference_loss = session.run(inference_loss_op, feed_dict={in_put_x: image,
                                                                           in_put_y: ground_truth})
                print(inference_loss)
                break


if __name__ == "__main__":
    main()
